from postanalyses.sleep_metrics import speedpair2wake, wake30_to_sleep30
import os
import json
from utils.plots import plot_daily
from utils.logs import log_daily
from utils.paths import is_contain_file
from postanalyses.trace_data import interp_by_lastcopy, get_contiguous_timeaxis, load_tracing_data
from postanalyses.motion import average_daily, bboxpair2centerpair, centerpair2speedpair, sumup
from postanalyses.sleep_metrics import speedpair2wake, wake30_to_sleep30
from VidSrc import create_mousesleep_video


def generate_videoseries_info(fp_info):
  root_chn = os.path.dirname(fp_info)
  if not os.path.exists(fp_info):
    # 构造一个VideoSeries
    vid_h = create_mousesleep_video(root_chn)
    with open(fp_info, 'w') as h:
        data_dict = {'nfr_list' : vid_h.nfr_list}
        json.dump(data_dict, h)


def recognize_sleep_for_channel(root_path, chn, offset):
  root_chn = os.path.join(root_path, str(chn))
  f_loc = 'locations.txt'
  f_info = str(chn) + '_info.json'
  fp_loc, fp_info = os.path.join(root_chn, f_loc), os.path.join(root_chn, f_info)
  if not os.path.exists(fp_info):
      generate_videoseries_info(fp_info)
  taxis, bodyhead_loc = load_tracing_data(fp_loc, fp_info)
  new_taxis = get_contiguous_timeaxis(taxis)
  new_bodyhead_loc = interp_by_lastcopy(taxis, bodyhead_loc)
  new_taxis.slice_by(stop=len(new_bodyhead_loc))
  centerpairs = bboxpair2centerpair(bodyhead_loc)
  speedpairs = centerpair2speedpair(centerpairs)
  wakes = speedpair2wake(speedpairs, window=60)
  wakes30 = sumup(wakes, 1800)
  # offset = 23 # offset=21 for dvr2
  # baseline sleep
  sleeps30_baseline = wake30_to_sleep30(wakes30[offset:offset+48+1])
  plot_daily(sleeps30_baseline[:48], 'before_inject_daily_sleep', 0.5, 'save', root_chn)
  log_daily(sleeps30_baseline[:48], 'before_inject_daily_sleep', 0.5, root_chn)
  # 1st inject
  wakes30_daily = wakes30[offset+48*2 : offset+48*3+1]
  sleeps30_daily = wake30_to_sleep30(wakes30_daily)
  plot_daily(sleeps30_daily[:48], 'after_1stinject_daily_sleep', 0.5, 'save', root_chn)
  log_daily(sleeps30_daily[:48], 'after_1stinject_daily_sleep', 0.5, root_chn)
  #wakes30_daily = average_daily(wakes30[48*4:48*6], offset=offset)
  wakes30_daily = wakes30[offset+48*4 : offset+48*5+1]
  sleeps30_daily = wake30_to_sleep30(wakes30_daily)
  plot_daily(sleeps30_daily[:48], 'after_2ndinject_daily_sleep', 0.5, 'save', root_chn)
  log_daily(sleeps30_daily[:48], 'after_2ndinject_daily_sleep', 0.5, root_chn)


if __name__ == '__main__':
  # LIGHT-ON 02:20, LIGHT-OFF 14:20
  root_path = '/home/youngway/data/mtrace/20201202_c57-males-dvr2'
  #offset = 23 # for dvr1
  offset=21 # for dvr2
  for subdir in os.listdir(root_path):
    dpath = os.path.join(root_path, subdir)
    if os.path.isdir(dpath) and is_contain_file(dpath, 'locations.txt'):
      print('processing channel %s' % dpath)
      recognize_sleep_for_channel(root_path, subdir, offset)
    
